import java.io.BufferedOutputStream;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.util.Arrays;

import opennlp.tools.doccat.DoccatModel;
import opennlp.tools.doccat.DocumentCategorizerEvaluator;
import opennlp.tools.doccat.DocumentCategorizerME;
import opennlp.tools.doccat.DocumentSample;
import opennlp.tools.doccat.DocumentSampleStream;
import opennlp.tools.util.InvalidFormatException;
import opennlp.tools.util.ObjectStream;
import opennlp.tools.util.PlainTextByLineStream;

public class ONLPClassfy {

    public static void main(String[] args) throws InvalidFormatException,
            IOException {
        new ONLPClassfy().train();
        //new ONLPClassfy().test2("BBVA ha firmado un acuerdo para vender su participación del 29,68% en Citic International Financial Holdings (CIFH), a la entidad China Citic Bank (CNCB) -matriz de CIFH-. El importe de la operación asciende a 845 millones de euros.");
        new ONLPClassfy().test("economía", "BBVA ha firmado un acuerdo para vender su participación del 29,68% en Citic International Financial Holdings (CIFH), a la entidad China Citic Bank (CNCB) -matriz de CIFH-. El importe de la operación asciende a 845 millones de euros.");
    }

    public void train() {
        String onlpModelPath = "/home/chema/model";
        String trainingDataFilePath = "/home/chema/training.txt";
        DoccatModel model = null;
        InputStream dataInputStream = null;
        OutputStream onlpModelOutput = null;
        try {
            dataInputStream = new FileInputStream(trainingDataFilePath);
            ObjectStream<String> lineStream = new PlainTextByLineStream(
                    dataInputStream, "UTF-8");
            ObjectStream<DocumentSample> sampleStream = new DocumentSampleStream(
                    lineStream);
            model = DocumentCategorizerME.train("es", sampleStream);
        } catch (IOException e) {
            System.err.println(e.getMessage());
        } finally {
            if (dataInputStream != null) {
                try {
                    dataInputStream.close();
                } catch (IOException e) {
                    System.err.println(e.getMessage());
                }
            }
        }
/*
* Now we are writing the calculated model to a file in order to use the
* trained classifier in production
*/
        try {
            onlpModelOutput = new BufferedOutputStream(new FileOutputStream(
                    onlpModelPath));
            model.serialize(onlpModelOutput);
        } catch (IOException e) {
            System.err.println(e.getMessage());
        } finally {
            if (onlpModelOutput != null) {
                try {
                    onlpModelOutput.close();
                } catch (IOException e) {
                    System.err.println(e.getMessage());
                }
            }
        }
    }

    /*
    * Now we call the saved model and test it
    * Give it a new text document and the expected category
    */
    public void test(String cat, String text) throws InvalidFormatException,
            IOException {
        String classificationModelFilePath = "/home/chema/model";
        InputStream is = new FileInputStream(classificationModelFilePath);
        DoccatModel classificationModel = new DoccatModel(is);
        DocumentCategorizerME classificationME = new DocumentCategorizerME(classificationModel);
        DocumentCategorizerEvaluator modelEvaluator = new DocumentCategorizerEvaluator(
                classificationME);
        String expectedDocumentCategory = cat;
        String documentContent = text;
        DocumentSample sample = new DocumentSample(expectedDocumentCategory,

                documentContent);
        double[] classDistribution = classificationME.categorize(documentContent);
        String predictedCategory = classificationME.getBestCategory(classDistribution);
        modelEvaluator.evaluteSample(sample);
        double result = modelEvaluator.getAccuracy();
        System.out.println("Model prediction : " + predictedCategory);
        System.out.println("Accuracy : " + result);
    }


    public void test2(String text) throws InvalidFormatException,
            IOException {
        String classificationModelFilePath = "/home/chema/model";
        InputStream is = new FileInputStream(classificationModelFilePath);
        DoccatModel classificationModel = new DoccatModel(is);
        DocumentCategorizerME classificationME = new DocumentCategorizerME(classificationModel);
        DocumentCategorizerEvaluator modelEvaluator = new DocumentCategorizerEvaluator(
                classificationME);
        String documentContent = text;
        double[] classDistribution = classificationME.categorize(documentContent);
        String predictedCategory = classificationME.getBestCategory(classDistribution);

        System.out.println("Model prediction : " + predictedCategory);
        System.out.println("------------------------------");
        Arrays.sort(classDistribution);
        for(int i=0;i<classDistribution.length;i++) {

            System.out.println(classificationME.getCategory(i)+" : "+classDistribution[i]);
        }
    }
}

